package com.nanbei.avl;

/**
 * @ClassDescription: AVL旋转树
 * @JdkVersion: 17
 * @Author: BaoShan.Li
 * @Date: 2024/6/4 22:31
 */
public class AVLTree {
    AVLNode root;

    public AVLTree(AVLNode root) {
        this.root = root;
    }

    /**
     * 向树中添加节点
     *
     * @param key 键
     * @param val 值
     */
    public void put(int key, Object val) {
        root = doPut(root, key, val);
    }


    private AVLNode doPut(AVLNode node, int key, Object val) {
        // 找到空位，创建节点
        if (node == null) {
            return new AVLNode(key, val);
        }
        // key已存在，更新val
        if (key == node.key) {
            node.val = val;
            return node;
        }
        if (key > node.key) {
            node.right = doPut(node.right, key, val);
        }
        if (key < node.key) {
            node.left = doPut(node.left, key, val);
        }
        // 更新节点的高度
        node.updateHeight(node);
        // 检查节点时候失衡然后调节失衡节点
        return node.balance(node);
    }

    /**
     * 删除节点
     *
     * @param key 需要删除的键
     */
    public void remove(int key) {
        root = doRemove(root, key);
    }

    private AVLNode doRemove(AVLNode node, int key) {
        if (node == null) {
            return null;
        }

        // 如果被删除的节点只有左孩子
        if (key > node.key) {
            node.right = doRemove(node.right, key);
        } else if (key < node.key) {
            node.left = doRemove(node.left, key);
        } else {
            if (node.right == null) {
                node = node.left;
            } else if (node.left == null) {
                node = node.right;
            } else {
                AVLNode right = node.right;
                while (right.left != null) {
                    right = right.left;
                }
                right.right = doRemove(node.right, right.key);
                right.left = node.left;
                node = right;
            }

        }
        node.updateHeight(node);
        return node.balance(node);
    }
}
